﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Net;

namespace EkspSys.Network
{
    /*
     * TODO: Remove magic numbers, and make code more robust instead of fixed buffer size.
     *       Lookup could maybe be optimized using binary search.
     */
    class ProcessHelper
    {
        private const int NO_ERROR = 0;

        // TODO: Fill out State type
        public enum State
        {
            CLOSED, LISTEN, SENT, RECIEVED, ESTABLISHED, FIN_WAIT_1
        };

        // TODO: Set state as State type.
        public struct TcpRow
        {
            public int state;
            public IPAddress localAddress;
            public UInt16 localPort;
            public IPAddress remoteAddr;
            public UInt16 remotePort;
            public Process owningProcess;
            public int offloadState;
        }

        public struct UdpRow
        {
            public IPAddress localAddress;
            public UInt16 localPort;
            public Process owningProcess;
        }

        [DllImport("iphlpapi.dll", SetLastError = true)]
        private static extern int GetTcpTable2(byte[] table, out int size, bool order);

        [DllImport("iphlpapi.dll", SetLastError = true)]
        private static extern int GetExtendedUdpTable(byte[] table, out int size, bool order, int af, int tableClass, int reserved);

        private static uint readDWORD(byte[] buf, ref int pos)
        {
            pos += 4;
            return (uint)(((uint)buf[pos - 1]) << 24) + (((uint)buf[pos - 2]) << 16) + (((uint)buf[pos - 3]) << 8) + ((uint)buf[pos - 4]);
        }

        private static ushort readPort(byte[] buf, ref int pos)
        {
            pos += 4;
            return (ushort)((((ushort)buf[pos - 4]) << 8) + ((ushort)buf[pos - 3]));
        }

        public static List<TcpRow> GetTcpConnections()
        {
            int bufSize = 20000;
            byte[] buf = new byte[bufSize];

            int res = GetTcpTable2(buf, out bufSize, true);
            if (res != NO_ERROR)
            {
                throw new ApplicationException("Failed loading TCP table. It properly is because the table buffer is too small.");
            }

            int pos = 0;
            uint entries = readDWORD(buf, ref pos);
            List<TcpRow> rows = new List<TcpRow>(Convert.ToInt32(entries));
            for (int i = 0; i < entries; i++)
            {
                TcpRow row = new TcpRow();
                row.state = Convert.ToInt32(readDWORD(buf, ref pos));
                row.localAddress = new IPAddress(readDWORD(buf, ref pos));
                row.localPort = readPort(buf, ref pos);
                row.remoteAddr = new IPAddress(readDWORD(buf, ref pos));
                row.remotePort = readPort(buf, ref pos);
                try
                {
                    row.owningProcess = Process.GetProcessById(Convert.ToInt32(readDWORD(buf, ref pos)));
                }
                catch (ArgumentException _)
                {
                    row.owningProcess = null;
                }
                pos += 4;

                rows.Add(row);
            }

            return rows;
        }

        public static List<UdpRow> GetUdpConnections()
        {
            int tableSize = 20000;
            byte[] table = new byte[tableSize];
            int res = GetExtendedUdpTable(table, out tableSize, true, 2, 1, 0);

            if (res != NO_ERROR)
            {
                throw new ApplicationException("Failed loading UDP table. It properly is because the table buffer is too small.");
            }

            int pos = 0;
            uint entries = readDWORD(table, ref pos);
            List<UdpRow> rows = new List<UdpRow>();
            for (int i = 0; i < entries; i++)
            {
                UdpRow row = new UdpRow();

                row.localAddress = new IPAddress(readDWORD(table, ref pos));
                row.localPort = readPort(table, ref pos);
                try
                {
                    row.owningProcess = Process.GetProcessById(Convert.ToInt32(readDWORD(table, ref pos)));
                }
                catch (ArgumentException _)
                {
                    row.owningProcess = null;
                }

                rows.Add(row);
            }
            return rows;
        }

        public static Process GetProcessForPort(Protocol proto, int lcalPrt)
        {
            if (proto == Protocol.TCP)
            {
                foreach (TcpRow row in GetTcpConnections())
                {
                    if (row.localPort == lcalPrt)
                        return row.owningProcess;
                }
                return null;
            }
            else if (proto == Protocol.UDP)
            {
                foreach (UdpRow row in GetUdpConnections())
                {
                    if (row.localPort == lcalPrt)
                        return row.owningProcess;
                }
                return null;
            }
            else
            {
                throw new ApplicationException("Unsupported protocol!");
            }
        }
    }
}
